# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

from experimental.prophet.graph import graph
from experimental.prophet.ops.util import (to_tuple,
                                           to_data_type_proto,
                                           to_duration_int,
                                           to_duration_str)
from experimental.prophet.contrib.functional import vectorize


@vectorize
def shift(value, window_size, auto_name=True):
  value = graph.as_variable(value)

  assert 0 < window_size and window_size < 1000000, window_size

  node = graph.create_node('timeseries.Shift', [value],
                           attrs={
                               'dtype': to_data_type_proto(value.dtype), 'window_size': window_size
                           })
  out_var = node.outputs[0]
  if auto_name and value.is_named_ref:
    out_var = out_var.named_ref(value.name + '.shift' + str(window_size), namespace='')

  return out_var


@vectorize
def time_shift(value, window_size, timestamp=None, auto_name=True):
  if timestamp is None:
    timestamp = graph.default_timestamp()

  value = graph.as_variable(value)
  window_size = to_duration_int(window_size)
  assert window_size >= 1000000, window_size  # 1ms

  node = graph.create_node('timeseries.TimeShift', [timestamp, value],
                           attrs={
                               'dtype': to_data_type_proto(value.dtype),
                               'window_size_ns': window_size
                           })

  # TODO(jaewon): named tuple, including 'timestamp' output
  out_var = node.outputs[1]
  if auto_name and value.is_named_ref:
    out_var = out_var.named_ref(value.name + '.' + to_duration_str(window_size), namespace='')

  return out_var


def time_moving_window(values, window_sizes, timestamp=None, auto_name=True):
  return _TimeMovingWindow(values, window_sizes, timestamp, auto_name)


class _TimeMovingWindow:
  def __init__(self, values, window_size, timestamp, auto_name):
    self._nodes = {}
    self._outputs = {}
    self._timestamp = timestamp if timestamp is not None else graph.default_timestamp()
    self._values = [graph.as_variable(v) for v in to_tuple(values)]
    self._window_sizes = [to_duration_int(ws) for ws in to_tuple(window_size)]
    self._auto_name = auto_name

  def _get_node(self, op_name):
    if op_name not in self._nodes:
      nodes = []
      for ws in self._window_sizes:
        assert ws >= 1000000, ws  # 1ms
        for v in self._values:
          node = graph.create_node(op_name, [self._timestamp, v],
                                   attrs={
                                       'dtype': to_data_type_proto(v.dtype), 'window_size_ns': ws
                                   })
          nodes.append((ws, v, node))
      self._nodes[op_name] = nodes
    return self._nodes[op_name]

  def _get_output(self, op_name, output_idx, name_postfix):
    key = (op_name, output_idx)
    if key in self._outputs:
      return self._outputs[key]

    nodes = self._get_node(op_name)
    if not self._auto_name:
      outputs = [n.outputs[output_idx] for ws, v, n in nodes]
      self._outputs[key] = outputs[0] if len(outputs) == 1 else outputs
      return self._outputs[key]

    outputs = []
    for ws, v, n in nodes:
      out_var = n.outputs[output_idx]
      if v.is_named_ref:
        out_var = out_var.named_ref(v.name + '.' + name_postfix + to_duration_str(ws), namespace='')
      outputs.append(out_var)

    self._outputs[key] = outputs[0] if len(outputs) == 1 else outputs
    return self._outputs[key]

  def sum(self):
    return self._get_output('timeseries.TimeMovingWindow_SumMeanOp', 0, 'msum')

  def mean(self):
    return self._get_output('timeseries.TimeMovingWindow_SumMeanOp', 1, 'mmean')

  def var(self):
    return self._get_output('timeseries.TimeMovingWindow_VarStdOp', 0, 'mvar')

  def std(self):
    return self._get_output('timeseries.TimeMovingWindow_VarStdOp', 1, 'mstd')

  def min(self):
    return self._get_output('timeseries.TimeMovingWindow_MinMaxOp', 0, 'mmin')

  def max(self):
    return self._get_output('timeseries.TimeMovingWindow_MinMaxOp', 1, 'mmax')

  def max_min_diff(self):
    return self._get_output('timeseries.TimeMovingWindow_MinMaxOp', 2, 'mmaxdiff')

  def twma(self):
    return self._get_output('timeseries.TimeWeightedMovingAverage', 0, 'twma')
